package com.guard.agent;

import cn.hutool.core.io.FileUtil;
import com.guard.agent.bridge.AgentContext;
import com.guard.agent.core.DynamicFieldAccessor;
import com.guard.agent.core.InjectBootstrapClassLoader;
import com.guard.agent.plugin.IPlugin;
import com.guard.agent.util.ClassLoaderInjectHelper;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.implementation.FieldAccessor;
import net.bytebuddy.jar.asm.Opcodes;
import net.bytebuddy.utility.JavaModule;
import org.tinylog.Logger;

import java.io.File;
import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.net.URLClassLoader;
import java.util.ServiceLoader;
import java.util.jar.JarFile;

import static net.bytebuddy.matcher.ElementMatchers.isSynthetic;
import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;

public class Bootstrap {

    private static final String DEBUG_PATH = "/Users/shangwq/Desktop/debug";

    public static void main(String args, Instrumentation inst) throws IOException {
        // 将guard-bridge.jar加入BootstrapClassLoader路径中
        inst.appendToBootstrapClassLoaderSearch(new JarFile(new File("/Users/shangwq/Documents/workspace/service-guard/guard-agent/bridge/target/guard-bridge.jar")));

        AgentContext.agentClassLoader = (URLClassLoader) Thread.currentThread().getContextClassLoader();

        // 注入Bootstrap ClassLoader
        ClassLoaderInjectHelper.inject(InjectBootstrapClassLoader.class, null, inst);

        AgentBuilder agentBuilder = new AgentBuilder.Default()
                .with(LISTENER)
                .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
                .with(AgentBuilder.TypeStrategy.Default.REDEFINE)
                .with(AgentBuilder.LocationStrategy.ForClassLoader.STRONG
                        .withFallbackTo(ClassFileLocator.ForClassLoader.ofSystemLoader()))
                .ignore(isSynthetic())
                .or(nameStartsWith("sun."))
                .or(nameStartsWith("com.sun."))
                .or(nameStartsWith("kotlin."))
                .or(nameStartsWith("javax."))
                .or(nameStartsWith("net.bytebuddy."))
                .or(nameStartsWith("com\\.sun\\.proxy\\.\\$Proxy.+"))
                .or(nameStartsWith("java\\.lang\\.invoke\\.BoundMethodHandle\\$Species_L.+"))
                .or(nameStartsWith("org.junit."))
                .or(nameStartsWith("junit."))
                .or(nameStartsWith("com.intellij."))
                .or(nameStartsWith("com.guard.agent"));

        ServiceLoader<IPlugin> pluginList = ServiceLoader.load(IPlugin.class, AgentContext.agentClassLoader);
        for (IPlugin plugin : pluginList) {
            System.out.println(plugin.namespace());
            agentBuilder = agentBuilder.type(plugin.typeMatcher())
                    .transform((builder, type, classLoader, module) -> {
                        if (plugin.isAddDynamicField()) {
                            builder = builder.defineField("guard$DynamicField", Object.class, Opcodes.ACC_PRIVATE | Opcodes.ACC_VOLATILE)
                                    .implement(DynamicFieldAccessor.class)
                                    .intercept(FieldAccessor.ofField("guard$DynamicField"));
                        }
                        builder = builder.visit(Advice.to(plugin.getAdvice()).on(plugin.methodMather()));
                        return builder;
                    });
        }

        agentBuilder.installOn(inst);
    }

    private static final AgentBuilder.Listener LISTENER = new AgentBuilder.Listener() {

        @Override
        public void onDiscovery(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
            // ignored
        }

        @Override
        public void onTransformation(TypeDescription td, ClassLoader ld, JavaModule m, boolean loaded, DynamicType dt) {
            Logger.info("onTransformation: {} from classLoader {} {}", td, ld, loaded);
            try {
                dt.saveIn(FileUtil.file(DEBUG_PATH));
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }

        @Override
        public void onIgnored(TypeDescription td, ClassLoader ld, JavaModule m, boolean loaded) {
            // ignored
        }

        @Override
        public void onError(String name, ClassLoader ld, JavaModule m, boolean loaded, Throwable error) {
            Logger.error(error, "onTransformation error name: {}  classLoader: {}", name, ld);
        }

        @Override
        public void onComplete(String name, ClassLoader ld, JavaModule m, boolean loaded) {
            // ignored
        }

    };

}
